{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 第二次课程作业——单视图重建"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 任务  \n",
    "\n",
    "- 根据所提供的符合单视重构要求的单张图片完成单视图重建任务。\n",
    "\n",
    "1. 加载使用 LABEL ME 工具标注产生的 json 数据，得到标记的平行线、平面交点、垂直面掩模。\n",
    "2. 据所标记的平行线在图像上的交点（通过计算得到），估算隐消点和场景中各个平面的隐消线。\n",
    "3. 根据由三组正交平行线得到的影消点，计算摄像机参数矩阵（标定）。\n",
    "4. 根据三个垂直面对应的隐消线和摄像机矩阵，计算场景平面的法方向（法向量）。\n",
    "5. 将平面交线的公共点代入具未知参数 d 的平面方程（设深度入=1），计算出各个场景平面到摄像机中心的距离 d。\n",
    "6. 将掩模区域中所有 2D 点坐标代入各点对应的场景平面方程（已知 d），计算出 3D 坐标，保存为 ply 文件，即单视图重建完成。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 预处理  \n",
    "\n",
    "- 以 images 文件里的 **chessboard.jpg** 图片为例，使用标注工具 labelme 对该图像进行标注，标记产生的 json 数据放在  **labelme_data** 文件夹中。首先在每个垂直面上标出若干条平行线，并确保三个平面上的平行线相互垂直。  \n",
    "<div align=center><img src=\"images/line.jpg\" width=200></div> \n",
    "\n",
    "- 同上，再次在三个平面上标注相互垂直的平行线，并与上一步所标注的平行线方向不同（示例图为垂直）。  \n",
    "<div align=center><img src=\"images/line'.jpg\" width=200></div>\n",
    "\n",
    "- 标出平面交线的公共点\n",
    "<div align=center><img src=\"images/intersection.jpg\" width=200></div>\n",
    "\n",
    "- 标出最终进行三维重建效果展示的区域（掩模）\n",
    "<div align=center><img src=\"images/plane.jpg\" width=200></div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 示例代码"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solution(A):  \n",
    "\n",
    "    \"\"\"Solves homogenous linear equations.\"\"\"\n",
    "\n",
    "    ## Ax = 0\n",
    "\n",
    "    # SVD\n",
    "    u, e, v = svd(A)\n",
    "    x = v.T[:, -1]  # min column of right singular vector\n",
    "\n",
    "    \"\"\" \n",
    "    ## Other Methods\n",
    "    # find the eigenvalues and eigenvector of A.T * A -- right singular \n",
    "    e_vals, e_vecs = np.linalg.eig(np.dot(A.T, A)) \n",
    "    # extract the eigenvector (column) associated with the minimum eigenvalue \n",
    "    x = np.array(e_vecs[:, np.argmin(np.abs(e_vals))])\n",
    "    \"\"\"\n",
    "\n",
    "    # np.set_printoptions(suppress=True)\n",
    "    # print(A, x, np.dot(A, x.reshape(4, 1)))\n",
    "\n",
    "    return x\n",
    "\n",
    "def get_label(img, data_line, data_plane):\n",
    "\n",
    "    \"\"\"Gets parallel lines and vertical planes from label json.\"\"\"\n",
    "\n",
    "    ## get parallel lines \n",
    "    \n",
    "    # print(data_line['shapes'])\n",
    "    line2points = defaultdict(list)  # store 2 points on each parallel line \n",
    "    line_names = [\"XZ_line\", \"YZ_line\", \"XY_line\", \"XZ_line'\", \"YZ_line'\", \"XY_line'\",]\n",
    "    intersect = []  # common point of plane intersection line\n",
    "    for each_shape in data_line['shapes']:\n",
    "        for each_line in line_names:\n",
    "            if each_shape['label'] == each_line:\n",
    "                # print(each_shape['points'])\n",
    "                # print([list(reversed(each_shape['points'][0])), list(reversed(each_shape['points'][1]))])\n",
    "                line2points[each_line].append([list(reversed(each_shape['points'][0])), list(reversed(each_shape['points'][1]))])\n",
    "            elif each_shape['label'] == \"intersection\":  # common point of plane intersection line\n",
    "                intersect = list(reversed(each_shape['points'][0])) + [1]\n",
    "    # print(line2points)\n",
    "\n",
    "    ## get vertical planes\n",
    "\n",
    "    lbl, lbl_names = utils.labelme_shapes_to_label(img.shape, data_plane['shapes']) \n",
    "\n",
    "    # lbl：binary array  in: 1  out: 0\n",
    "    # lbl_names ：dict   _background_：0   other labels: positive\n",
    "    # print(lbl_names)\n",
    "    mask = []  #  for 3 verticle planes\n",
    "    for i in range(1, len(lbl_names)):  # ignore background\n",
    "        mask.append((lbl == i).astype(np.uint8))  \n",
    "\n",
    "    # captions = ['%d: %s' % (l, name) for l, name in enumerate(lbl_names)]\n",
    "    # lbl_viz = utils.draw_label(lbl, img, captions)\n",
    "    # plt.imshow(lbl_viz)\n",
    "    # plt.show()\n",
    "\n",
    "    return line2points, intersect, mask\n",
    "\n",
    "def cal_vanish(line2points):\n",
    "\n",
    "    \"\"\"Calculates vanishing points and horizon lines.\"\"\"\n",
    "    \n",
    "    ## calculate vanishing points\n",
    "    vanish_points = defaultdict(list)  # for pairs of parallel lines \n",
    "    for each_line, points in line2points.items():  # dict\n",
    "        # print(each_line, len(points))\n",
    "        line = []  # for each pair of parallel lines, store several line equations\n",
    "        for each_point in points:  # points array\n",
    "            # cross product -- calculate line equation\n",
    "            # line.append(np.cross(each_point[0] + [1], each_point[1] + [1]))  # homogenous coordinates\n",
    "            line.append(solution(np.array([each_point[0] + [1], each_point[1] + [1]])))\n",
    "        # print(line)\n",
    "        # p = np.cross(line[0], line[1])  # homogenous -- up to a scale\n",
    "        p = solution(np.array(line))    \n",
    "        vanish_points[each_line] = p / p[-1]\n",
    "    # print(vanish_points)\n",
    "\n",
    "    ## calculate vanishing lines\n",
    "    \"\"\"\n",
    "    horizs = np.array([np.cross(vanish_points[\"XZ_line\"], vanish_points[\"XZ_line'\"]), \n",
    "                       np.cross(vanish_points[\"YZ_line\"], vanish_points[\"YZ_line'\"]),\n",
    "                       np.cross(vanish_points[\"XY_line\"], vanish_points[\"XY_line'\"])]).T  \n",
    "    \"\"\"\n",
    "    horizs = np.array([solution(np.array([vanish_points[\"XZ_line\"], vanish_points[\"XZ_line'\"]])), \n",
    "                       solution(np.array([vanish_points[\"YZ_line\"], vanish_points[\"YZ_line'\"]])), \n",
    "                       solution(np.array([vanish_points[\"XY_line\"], vanish_points[\"XY_line'\"]]))]).T\n",
    "    # print(horizs)\n",
    "\n",
    "    return vanish_points, horizs\n",
    "\n",
    "def calibrate(vanish_points):\n",
    "    \n",
    "    \"\"\"Calibrates the camera (square pixels & no skew).\"\"\" \n",
    "\n",
    "    # W = [[w1 w2 w4] [w2 w3 w5] [w4 w5 w6]]\n",
    "    # (square pixels: w2 = 0  no skew: w1 = w3)\n",
    "\n",
    "    # θ = 90° -- v1.T * W * v2 = 0 ; v2.T * W * v3 = 0 ; v1.T * W * v4 = 0\n",
    "    # (a1a2 + b1b2)w1 + (c1a2 + a1c2)w4 + (c1b2 + b1c2)w5 + c1c2w6 = 0\n",
    "    # (a1a3 + b1b3)w1 + (c1a3 + a1c3)w4 + (c1b3 + b1c3)w5 + c1c3w6 = 0\n",
    "    # (a2a3 + b2b3)w1 + (c2a3 + a2c3)w4 + (c2b3 + b2c3)w5 + c2c3w6 = 0      \n",
    "\n",
    "    a1, b1, c1 = vanish_points[\"XZ_line\"]\n",
    "    a2, b2, c2 = vanish_points[\"YZ_line\"]\n",
    "    a3, b3, c3 = vanish_points[\"XY_line\"]\n",
    "\n",
    "    param = np.array([[a1 * a2 + b1 * b2, c1 * a2 + a1 * c2, c1 * b2 + b1 * c2, c1 * c2], \n",
    "                      [a1 * a3 + b1 * b3, c1 * a3 + a1 * c3, c1 * b3 + b1 * c3, c1 * c3],\n",
    "                      [a2 * a3 + b2 * b3, c2 * a3 + a2 * c3, c2 * b3 + b2 * c3, c2 * c3]])  # parameter matrix  \n",
    "\n",
    "    # print(np.linalg.matrix_rank(param))  # rank < 4 -- non-zero solution\n",
    "    \n",
    "    w1, w4, w5, w6 = solution(param)  # solve w1, w4, w5, w6 -- up to a scale\n",
    "    W = np.array([[w1, 0, w4], [0, w1, w5], [w4, w5, w6]])\n",
    "    W /= W[-1, -1]\n",
    "    \n",
    "    ### Once W is calculated, we get camera matrix K: \n",
    "\n",
    "    ## Cholesky factorization -- W = (K.inv).T * (K.inv)\n",
    "    # e_vals, e_vecs = np.linalg.eig(W) \n",
    "    # print(e_vals)  # all e_vals > 0 -- positive denifite\n",
    "    L = cholesky(W, lower=False)  # W = L.T * L, L is upper trianguler matrix \n",
    "    K = np.linalg.inv(L)  # K is upper trianguler matrix \n",
    "    K /= K[-1][-1]\n",
    "    # print(K)\n",
    "\n",
    "    \"\"\"\n",
    "    ### Other Methods\n",
    "    ## directly -- W = (K * K.T).inv\n",
    "    # K = [[f 0 u0] [0 f v0] [0 0 1]] (square pixels & no skew)\n",
    "    # K * K.T = [[f^2 + u0^2 u0v0 u0] [u0v0 f^2 + v0^2 v0] [u0 v0 1]]\n",
    "    W_inv = np.linalg.inv(W) \n",
    "    dot = W_inv / W_inv[-1][-1]\n",
    "    # print(W, W_inv, dot)\n",
    "    u0 = dot[-1][0]\n",
    "    v0 = dot[-1][1]\n",
    "    # assert u0 * v0 == dot[-1][0] * dot[-1][1]\n",
    "    # assert abs((u0 * u0 - v0 * v0) - (dot[0][0] - dot[1][1])) < 1e-3\n",
    "    f = np.sqrt(dot[1][1] - v0 * v0)\n",
    "    K = np.array([[f, 0, u0], [0, f, v0], [0, 0, 1]])\n",
    "    # print(K)\n",
    "    \"\"\"\n",
    "    \n",
    "    return K\n",
    "\n",
    "def cal_3D(K, p_2D_H):\n",
    "\n",
    "    \"\"\"Calculates unit 3D pos without taking projective depth into account.\"\"\"\n",
    "\n",
    "    p_3D = np.dot(np.linalg.inv(K), p_2D_H)\n",
    "    p_3D /= np.linalg.norm(p_3D)  # p / || p ||\n",
    "\n",
    "    return p_3D\n",
    "\n",
    "def cal_scene(K, horizs, intersect):\n",
    "\n",
    "    \"\"\"Calculates scene plane equations.\"\"\"\n",
    "\n",
    "    # X_H.T * Pi = 0 (homogenous) -- Pi = (unit_n.T, d).T\n",
    "    # unit_n.T * X + d = 0\n",
    "\n",
    "    ## calculate scene plane orientations -- normal vector\n",
    "    \n",
    "    # n = K.T * l_horiz\n",
    "    N = np.dot(K.T, horizs)\n",
    "    # unit_n = n / || n ||\n",
    "    unit_N = N / np.linalg.norm(N, axis=0)\n",
    "    # print(unit_N)\n",
    "\n",
    "    ## calculate distance between plane and camera center\n",
    "\n",
    "    # common point of plane intersection lines in 3D \n",
    "    # X_H = (λ * K.inv * x / || K.inv * x ||, 1) -- suppose projective depth λ is 1\n",
    "    intersect_3D = cal_3D(K, intersect).reshape((3,1))\n",
    "    # print(intersect_3D)\n",
    "    # intersect_3D_H = np.append(intersect_3D, 1)\n",
    "\n",
    "    # substitute the common point into plane equations\n",
    "    D = -1 * np.dot(unit_N.T, intersect_3D)\n",
    "    # print(D)\n",
    "\n",
    "    # unit_N.T * X + D = 0 -- X_H * Pi = 0\n",
    "    Pi = np.concatenate((unit_N, D.T), axis=0)  # (4, 3)\n",
    "\n",
    "    return Pi\n",
    "\n",
    "def reconstruction(K, Pi, img, masks):\n",
    "\n",
    "    \"\"\"Reconstructs 3D points in masked image for 3 verticle planes.\"\"\"\n",
    "\n",
    "    assert img.shape[:-1] == np.array(masks).shape[1:]\n",
    "\n",
    "    pos = []  # position in 3D\n",
    "    rgb = []  # pixel in original image\n",
    "    for axis, each_mask in enumerate(masks):  # plane XZ; YZ; XY\n",
    "        for i in range(img.shape[0]):\n",
    "            for j in range(img.shape[1]):\n",
    "                if each_mask[i, j] == 1:\n",
    "                    # X_H * Pi = 0 -- X * Pi[:-1] + Pi[-1] = 0\n",
    "                    # X_H = (λ * K.inv * x / || K.inv * x ||, 1) -- solve λ\n",
    "                    point = cal_3D(K, [i, j, 1]).reshape((3,1))\n",
    "                    lambd = -Pi[-1, axis] / np.dot(Pi[:-1, axis], point)\n",
    "                    pos.append(lambd * point)\n",
    "                    rgb.append(img[i, j])\n",
    "\n",
    "    return pos, rgb\n",
    "\n",
    "def create_output(vertices, colors, filename):\n",
    "\n",
    "    \"\"\"Creates point cloud file.\"\"\"\n",
    "\n",
    "    vertices = np.hstack([np.array(vertices).reshape(-1, 3), np.array(colors).reshape(-1, 3)])\n",
    "    np.savetxt(filename, vertices, fmt='%f %f %f %d %d %d') \n",
    "    ply_header = '''ply\\nformat ascii 1.0\\nelement vertex %(vert_num)d\\nproperty float x\\nproperty float y\\nproperty float z\\nproperty uchar red\\nproperty uchar green\\nproperty uchar blue\\nend_header\\n'''\n",
    "    with open(filename, 'r+') as f:\n",
    "        old = f.read()\n",
    "        f.seek(0)\n",
    "        f.write(ply_header % dict(vert_num=len(vertices)))\n",
    "        f.write(old)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 结果测试"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **导入相关模块**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import matplotlib.image as mpimg\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from labelme import utils\n",
    "from collections import defaultdict\n",
    "from scipy.linalg import svd, cholesky, qr\n",
    "import open3d as o3d\n",
    "from open3d import web_visualizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **加载 json 文件中的标记数据**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load data\n",
    "\n",
    "json_file_line = r'labelme_data\\chessboard_line.json'  # label the parallel lines\n",
    "json_file_plane = r'labelme_data\\chessboard_plane.json'  # label the vertical planes\n",
    "data_line = json.load(open(json_file_line))\n",
    "data_plane = json.load(open(json_file_plane))\n",
    "# img = utils.img_b64_to_arr(data_plane['imageData'])  # png: 4 channels\n",
    "img = mpimg.imread(r'images\\chessboard.jpg')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **单视图标定与重构，并保存 3D 模型**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[WARNING] shape:labelme_shapes_to_label:79 - labelme_shapes_to_label is deprecated, so please use shapes_to_label.\n"
     ]
    }
   ],
   "source": [
    "line2points, intersect, masks = get_label(img, data_line, data_plane)\n",
    "\n",
    "# calculate points and lines at infinity from labeled parallel lines\n",
    "\n",
    "vanish_points, vanish_lines = cal_vanish(line2points)\n",
    "\n",
    "# according to intersection points of 3 pairs parallel lines which are mutually orthogonal,\n",
    "# calculate camera matrix\n",
    "\n",
    "K = calibrate(vanish_points)\n",
    "\n",
    "# according to lines at infinity of 3 vertical planes and camera matrix, \n",
    "# calculate scene plane orientations (normal vectors);\n",
    "# substitute the common point(assuming projective depth) of plane intersection lines into plane equations, \n",
    "# calculate distance between each plane and camera center\n",
    "\n",
    "Pi = cal_scene(K, vanish_lines, intersect)\n",
    "\n",
    "# substituting 2D points into corresponding plane equation,\n",
    "# calculate 3D positions for masked image (up to a unknown scale)\n",
    "\n",
    "pos, rgb = reconstruction(K, Pi, img, masks)\n",
    "\n",
    "# save as ply file\n",
    "\n",
    "create_output(pos, rgb, r'chessboard_3D.ply')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **可视化 3D 模型**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5aaa2b66541d464d84dbd60dcf73d768",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "WebVisualizer(window_uid='window_1')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# visualization of point clouds.\n",
    "pcd = o3d.io.read_point_cloud('chessboard_3D.ply')\n",
    "# o3d.visualization.draw_geometries([pcd])\n",
    "web_visualizer.draw(pcd)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "2fd6ff00ff4a419d324d5b7e4b1b0789b8ee895e93e15d812a34184c59464f6c"
  },
  "kernelspec": {
   "display_name": "Python 3.8.11 64-bit ('MyCV': conda)",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
